package zkp

import (
	"bytes"
	"encoding/base64"
	"encoding/hex"
	"fmt"
	"math/big"

	"github.com/consensys/gnark-crypto/ecc"
	"github.com/consensys/gnark-crypto/ecc/bn254/fr/mimc"
	"github.com/consensys/gnark/backend/groth16"
	"github.com/consensys/gnark/constraint"
	"github.com/consensys/gnark/frontend"
)

// Prover ZKP 证明生成器
type Prover struct {
	provingKey groth16.ProvingKey
	ccs        constraint.ConstraintSystem
}

// NewProver 创建新的证明生成器
func NewProver(provingKey groth16.ProvingKey, ccs constraint.ConstraintSystem) *Prover {
	return &Prover{
		provingKey: provingKey,
		ccs:        ccs,
	}
}

// ComputeMiMCHash 计算 MiMC 哈希 (用于计算 commitment 和 response)
// 注意: gnark-crypto 的 MiMC 实现与电路中的 gnark/std/hash/mimc 应该产生相同结果
// 电路中使用: mimcHasher.Write(a); mimcHasher.Write(b) - 按顺序处理两个 field elements
//
// ⚠️ 重要: 所有输入必须先模 r（BN254 scalar field modulus），与电路保持一致
func ComputeMiMCHash(secretHex string, dataBytes []byte) (string, error) {
	// BN254 scalar field 模数
	modulus, _ := new(big.Int).SetString("21888242871839275222246405745257275088548364400416034343698204186575808495617", 10)

	// ✅ Secret 是 hex 字符串，需要解码为字节
	secretBytes, err := hex.DecodeString(secretHex)
	if err != nil {
		return "", fmt.Errorf("解析 secret hex 失败: %v", err)
	}

	// ✅ 关键修复：将 secret 模 r，与电路中的行为一致
	secretBig := new(big.Int).SetBytes(secretBytes)
	secretBig.Mod(secretBig, modulus)
	secretFieldBytes := make([]byte, 32)
	secretBig.FillBytes(secretFieldBytes)

	// ✅ 关键修复：将 data 模 r，与电路中的行为一致
	dataBig := new(big.Int).SetBytes(dataBytes)
	dataBig.Mod(dataBig, modulus)
	dataFieldBytes := make([]byte, 32)
	dataBig.FillBytes(dataFieldBytes)

	// 创建 MiMC 哈希
	// gnark-crypto 的 MiMC 实现: 按顺序写入两个域元素
	h := mimc.NewMiMC()
	h.Write(secretFieldBytes)  // 第一个域元素（32 字节，已模 r）
	h.Write(dataFieldBytes)    // 第二个域元素（32 字节，已模 r）

	// 返回十六进制结果
	result := h.Sum(nil)
	return hex.EncodeToString(result), nil
}

// GenerateProof 生成零知识证明
// 参数:
//   - secretHex: 设备密钥 (十六进制字符串，64字符=32字节)
//   - deviceID: 设备 ID
//   - challenge: 挑战值 (十六进制字符串)
//   - commitment: 预先计算的承诺值 (十六进制字符串)
//
// 返回:
//   - proofBase64: Base64 编码的证明
//   - deviceIDHex: deviceID 的十六进制表示（用于 public_witness）
//   - response: MiMC(secret, challenge) 的十六进制结果
func (p *Prover) GenerateProof(secretHex, deviceID, challenge, commitment string) (proofBase64, deviceIDHex, response string, err error) {
	// 1. 解析输入数据
	// ✅ Secret 是 hex 字符串，需要解码
	secretBytes, err := hex.DecodeString(secretHex)
	if err != nil {
		return "", "", "", fmt.Errorf("解析 secret 失败: %v", err)
	}

	challengeBytes, err := hex.DecodeString(challenge)
	if err != nil {
		return "", "", "", fmt.Errorf("解析 challenge 失败: %v", err)
	}

	commitmentBytes, err := hex.DecodeString(commitment)
	if err != nil {
		return "", "", "", fmt.Errorf("解析 commitment 失败: %v", err)
	}

	// ✅ 修复：将 deviceID 转换为域元素（32 字节）用于 public_witness
	deviceIDBytes := []byte(deviceID)
	deviceIDBig := new(big.Int).SetBytes(deviceIDBytes)
	deviceIDFieldBytes := make([]byte, 32)
	deviceIDBig.FillBytes(deviceIDFieldBytes)
	deviceIDHex = hex.EncodeToString(deviceIDFieldBytes)

	// 2. 计算 response = MiMC(secret, challenge)
	// ✅ 修复：使用 ComputeMiMCHash 确保与电路一致
	response, err = ComputeMiMCHash(secretHex, challengeBytes)
	if err != nil {
		return "", "", "", fmt.Errorf("计算 response 失败: %v", err)
	}
	responseBytes, _ := hex.DecodeString(response)

	// 3. 转换为 big.Int (用于电路)
	// ✅ 关键修复：所有值都必须模 r，与电路行为一致
	modulus, _ := new(big.Int).SetString("21888242871839275222246405745257275088548364400416034343698204186575808495617", 10)

	secretBig := new(big.Int).SetBytes(secretBytes)
	secretBig.Mod(secretBig, modulus)

	challengeBig := new(big.Int).SetBytes(challengeBytes)
	challengeBig.Mod(challengeBig, modulus)

	commitmentBig := new(big.Int).SetBytes(commitmentBytes)
	commitmentBig.Mod(commitmentBig, modulus)

	responseBig := new(big.Int).SetBytes(responseBytes)
	responseBig.Mod(responseBig, modulus)

	// 4. 创建见证 (witness)
	assignment := &AuthCircuit{
		Secret:     secretBig,
		DeviceID:   deviceIDBig,
		Challenge:  challengeBig,
		Commitment: commitmentBig,
		Response:   responseBig,
	}

	// 5. 生成完整见证
	witness, err := frontend.NewWitness(assignment, ecc.BN254.ScalarField())
	if err != nil {
		return "", "", "", fmt.Errorf("创建见证失败: %v", err)
	}

	// 6. 生成 Groth16 证明
	proof, err := groth16.Prove(p.ccs, p.provingKey, witness)
	if err != nil {
		return "", "", "", fmt.Errorf("生成证明失败: %v", err)
	}

	// 7. 序列化证明为字节
	var buf bytes.Buffer
	if _, err := proof.WriteTo(&buf); err != nil {
		return "", "", "", fmt.Errorf("序列化证明失败: %v", err)
	}

	// 8. Base64 编码
	proofBase64 = base64.StdEncoding.EncodeToString(buf.Bytes())

	return proofBase64, deviceIDHex, response, nil
}

// 注意: 客户端不需要 SetupKeys() 函数
// Trusted Setup 应该由开发团队在可信环境中执行一次
// 生成的 auth_proving.key 分发给所有客户端
// 生成的 auth_verifying.key 部署到 Edge 服务端
